import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, Any

def plot_backtest_results(results: Dict[str, Any]):
    """
    可视化回测结果
    
    Args:
        results: 回测结果字典
    """
    # 设置绘图风格
    plt.style.use('seaborn')
    sns.set_palette("husl")
    
    # 创建子图
    fig = plt.figure(figsize=(15, 10))
    gs = fig.add_gridspec(3, 2)
    
    # 1. 资金曲线
    ax1 = fig.add_subplot(gs[0, :])
    holdings = results['Holdings History']
    holdings['equity'].plot(ax=ax1, label='账户净值')
    ax1.set_title('账户净值曲线')
    ax1.set_xlabel('时间')
    ax1.set_ylabel('净值 ($)')
    ax1.legend()
    ax1.grid(True)
    
    # 2. 回撤曲线
    ax2 = fig.add_subplot(gs[1, 0])
    drawdown = holdings['drawdown']
    drawdown.plot(ax=ax2, color='red', alpha=0.5)
    ax2.fill_between(drawdown.index, drawdown.values, alpha=0.3, color='red')
    ax2.set_title('回撤曲线')
    ax2.set_xlabel('时间')
    ax2.set_ylabel('回撤 (%)')
    ax2.grid(True)
    
    # 3. 收益分布
    ax3 = fig.add_subplot(gs[1, 1])
    returns = holdings['returns']
    sns.histplot(returns, kde=True, ax=ax3)
    ax3.set_title('收益分布')
    ax3.set_xlabel('收益率')
    ax3.set_ylabel('频率')
    
    # 4. 交易统计
    ax4 = fig.add_subplot(gs[2, 0])
    trade_stats = pd.Series({
        '总交易次数': results['Total Trades'],
        '盈利交易': results['Winning Trades'],
        '亏损交易': results['Losing Trades']
    })
    trade_stats.plot(kind='bar', ax=ax4)
    ax4.set_title('交易统计')
    ax4.set_ylabel('次数')
    
    # 5. 关键指标
    ax5 = fig.add_subplot(gs[2, 1])
    metrics = pd.Series({
        '夏普比率': results['Sharpe Ratio'],
        '年化收益率': results['Annual Return'],
        '最大回撤': results['Max Drawdown'],
        '胜率': results['Win Rate']
    })
    metrics.plot(kind='bar', ax=ax5)
    ax5.set_title('关键指标')
    ax5.set_ylabel('数值')
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图表
    plt.savefig('backtest_results.png')
    print("回测结果图表已保存到 backtest_results.png")
    
    # 显示图表
    plt.show()

def plot_ml_feature_importance(feature_importance: pd.Series):
    """
    可视化机器学习特征重要性
    
    Args:
        feature_importance: 特征重要性Series
    """
    plt.figure(figsize=(12, 6))
    feature_importance.sort_values(ascending=True).plot(kind='barh')
    plt.title('特征重要性')
    plt.xlabel('重要性得分')
    plt.tight_layout()
    plt.savefig('feature_importance.png')
    print("特征重要性图表已保存到 feature_importance.png")
    plt.show()

def plot_trading_signals(market_data: pd.DataFrame, signals: pd.DataFrame):
    """
    可视化交易信号
    
    Args:
        market_data: 市场数据
        signals: 交易信号数据
    """
    plt.figure(figsize=(15, 7))
    
    # 绘制价格
    plt.plot(market_data.index, market_data['close'], label='价格', alpha=0.7)
    
    # 绘制买入信号
    buy_signals = signals[signals['signal'] == 1]
    plt.scatter(buy_signals.index, market_data.loc[buy_signals.index, 'close'],
                marker='^', color='g', s=100, label='买入信号')
    
    # 绘制卖出信号
    sell_signals = signals[signals['signal'] == -1]
    plt.scatter(sell_signals.index, market_data.loc[sell_signals.index, 'close'],
                marker='v', color='r', s=100, label='卖出信号')
    
    plt.title('交易信号')
    plt.xlabel('时间')
    plt.ylabel('价格')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    
    plt.savefig('trading_signals.png')
    print("交易信号图表已保存到 trading_signals.png")
    plt.show()

def plot_portfolio_analysis(results: Dict[str, Any]):
    """
    绘制投资组合分析图表
    
    Args:
        results: 回测结果字典
    """
    holdings = results['Holdings History']
    
    # 创建图表
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    
    # 1. 滚动收益率
    rolling_returns = holdings['returns'].rolling(window=30).mean()
    axes[0, 0].plot(rolling_returns.index, rolling_returns.values)
    axes[0, 0].set_title('30日滚动收益率')
    axes[0, 0].grid(True)
    
    # 2. 滚动波动率
    rolling_vol = holdings['returns'].rolling(window=30).std() * np.sqrt(252)
    axes[0, 1].plot(rolling_vol.index, rolling_vol.values)
    axes[0, 1].set_title('30日滚动波动率')
    axes[0, 1].grid(True)
    
    # 3. 滚动夏普比率
    risk_free_rate = 0.02  # 假设无风险利率为2%
    excess_returns = holdings['returns'] - risk_free_rate/252
    rolling_sharpe = (excess_returns.rolling(window=30).mean() / 
                     holdings['returns'].rolling(window=30).std() * 
                     np.sqrt(252))
    axes[1, 0].plot(rolling_sharpe.index, rolling_sharpe.values)
    axes[1, 0].set_title('30日滚动夏普比率')
    axes[1, 0].grid(True)
    
    # 4. 收益率QQ图
    from scipy import stats
    returns = holdings['returns'].dropna()
    stats.probplot(returns, dist="norm", plot=axes[1, 1])
    axes[1, 1].set_title('收益率QQ图')
    
    plt.tight_layout()
    plt.savefig('portfolio_analysis.png')
    print("投资组合分析图表已保存到 portfolio_analysis.png")
    plt.show()

def plot_trading_results(equity_curve: pd.DataFrame, trades: pd.DataFrame, save_path: str = None):
    """
    绘制交易结果图表
    
    Args:
        equity_curve: 资金曲线数据
        trades: 交易记录数据
        save_path: 保存图表的路径（可选）
    """
    # 设置图表风格
    plt.style.use('seaborn')
    fig = plt.figure(figsize=(15, 10))
    
    # 1. 资金曲线
    ax1 = plt.subplot(2, 2, 1)
    equity_curve['total'].plot(ax=ax1, label='Portfolio Value')
    ax1.set_title('Portfolio Equity Curve')
    ax1.set_xlabel('Date')
    ax1.set_ylabel('Portfolio Value ($)')
    ax1.legend()
    
    # 2. 回撤图
    ax2 = plt.subplot(2, 2, 2)
    equity_curve['drawdown'].plot(ax=ax2, label='Drawdown', color='red')
    ax2.set_title('Portfolio Drawdown')
    ax2.set_xlabel('Date')
    ax2.set_ylabel('Drawdown (%)')
    ax2.legend()
    
    # 3. 收益分布直方图
    if not trades.empty:
        ax3 = plt.subplot(2, 2, 3)
        sns.histplot(data=trades['pnl'], bins=30, ax=ax3)
        ax3.set_title('PnL Distribution')
        ax3.set_xlabel('Profit/Loss ($)')
        ax3.set_ylabel('Frequency')
        
        # 4. 累计收益曲线
        ax4 = plt.subplot(2, 2, 4)
        trades['cumulative_pnl'] = trades['pnl'].cumsum()
        trades['cumulative_pnl'].plot(ax=ax4, label='Cumulative PnL')
        ax4.set_title('Cumulative PnL')
        ax4.set_xlabel('Trade Number')
        ax4.set_ylabel('Cumulative Profit/Loss ($)')
        ax4.legend()
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
    else:
        plt.show()
    
    plt.close()

def plot_trade_analysis(trades: pd.DataFrame, save_path: str = None):
    """
    绘制交易分析图表
    
    Args:
        trades: 交易记录数据
        save_path: 保存图表的路径（可选）
    """
    if trades.empty:
        print("No trades to analyze")
        return
        
    plt.style.use('seaborn')
    fig = plt.figure(figsize=(15, 10))
    
    # 1. 每月收益热力图
    trades['month'] = pd.to_datetime(trades['timestamp']).dt.month
    trades['year'] = pd.to_datetime(trades['timestamp']).dt.year
    monthly_returns = trades.pivot_table(
        values='pnl', 
        index='year',
        columns='month',
        aggfunc='sum'
    )
    
    ax1 = plt.subplot(2, 2, 1)
    sns.heatmap(monthly_returns, annot=True, fmt='.0f', cmap='RdYlGn', ax=ax1)
    ax1.set_title('Monthly Returns Heatmap')
    
    # 2. 胜率饼图
    ax2 = plt.subplot(2, 2, 2)
    win_count = len(trades[trades['pnl'] > 0])
    loss_count = len(trades[trades['pnl'] < 0])
    plt.pie([win_count, loss_count], labels=['Wins', 'Losses'], autopct='%1.1f%%')
    ax2.set_title('Win/Loss Ratio')
    
    # 3. 交易量分析
    ax3 = plt.subplot(2, 2, 3)
    trades['hour'] = pd.to_datetime(trades['timestamp']).dt.hour
    trades.groupby('hour')['pnl'].count().plot(kind='bar', ax=ax3)
    ax3.set_title('Trading Volume by Hour')
    ax3.set_xlabel('Hour of Day')
    ax3.set_ylabel('Number of Trades')
    
    # 4. 收益率箱线图
    ax4 = plt.subplot(2, 2, 4)
    trades.boxplot(column='pnl', by='direction', ax=ax4)
    ax4.set_title('PnL Distribution by Direction')
    ax4.set_ylabel('Profit/Loss ($)')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
    else:
        plt.show()
    
    plt.close()
